import torchvision.datasets as dsets
import torch
import torchvision.transforms as transform
import matplotlib.pyplot as plt

batch_size=64

def datasetting():
    # 设置数据集
    train_data = dsets.MNIST(root='data', train=True, transform=transform.ToTensor(), download=True)
    test_data = dsets.MNIST(root='data', train=False, transform=transform.ToTensor(), download=True)

    # 读取数据 分batch
    train_loader = torch.utils.data.DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True,num_workers=8)

    # 把test_data分成校验集和测试集
    indices = range(len(test_data))
    indices_val = indices[:5000]
    indices_test = indices[5000:]

    sampler_val = torch.utils.data.sampler.SubsetRandomSampler(indices_val)
    sampler_test = torch.utils.data.sampler.SubsetRandomSampler(indices_test)

    val_loader = torch.utils.data.DataLoader(dataset=test_data, batch_size=batch_size, shuffle=False,
                                             sampler=sampler_val,num_workers=8)
    test_loader = torch.utils.data.DataLoader(dataset=test_data, batch_size=batch_size, shuffle=False,
                                              sampler=sampler_test,num_workers=8)






    return train_loader, val_loader, test_loader


